#!/usr/bin/env python3
# coding: utf-8
# Copyright 2024 Huawei Technologies Co., Ltd
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
# ===========================================================================
import os
import json
import platform

from ansible.module_utils.basic import AnsibleModule
from ansible.module_utils import common_info, common_utils, venv_installer
from ansible.module_utils.common_info import DeployStatus


class NpuExporterInstaller:
    def __init__(self):
        self.module = AnsibleModule(
            argument_spec=dict(
                tags=dict(type="list", required=True),
                resources_dir=dict(type="str", required=True),
                ip_address=dict(type="str", required=True),
            )
        )
        self.run_tags = self.module.params["tags"]
        self.username = "hwMindX"
        self.uid = "9000"
        self.log_path = "/var/log/mindx-dl/npu-exporter"
        self.install_path = "/usr/local/bin/npu-exporter"
        self.arch = platform.machine()
        self.resources_dir = os.path.expanduser(self.module.params["resources_dir"])
        self.ip_address = self.module.params["ip_address"]
        self.messages = []

    def _module_failed(self):
        return self.module.fail_json(msg="\n".join(self.messages), rc=1, changed=False)

    def _module_success(self):
        return self.module.exit_json(msg="Install npu_exporter success.", rc=0, changed=True)

    def _find_npu_exporter_pkg(self, pattern):
        pkgs, msgs = common_utils.find_files(os.path.join(self.resources_dir, "NpuExporter*"), pattern)
        self.messages.extend(msgs)
        if not pkgs:
            if "auto" in self.run_tags:
                self.module.exit_json(std_out="can not find npu_exporter package, npu_exporter install skipped", rc=0,
                                      result={DeployStatus.DEPLOY_STATUS: DeployStatus.SKIP},
                                      changed=False)
            else:
                self.messages.append("can not find npu_exporter package.")
                self._module_failed()

        return pkgs[0]
    
    def _get_config(self):
        config_path = self._find_npu_exporter_pkg("npu_monitor_config.json")
        with open(config_path, 'r') as f:
            self.config = json.load(f).get("npu_exporter", {})
            self.port = self.config.get("port", 8082)

    def create_user(self):
        rc, _, _ = self.module.run_command(["id", self.username])
        if rc != 0:
            self.module.run_command(["useradd", "-d", "/home{}".format(self.username),
                "-u", self.uid, "-m", "-s", "/sbin/nologin", self.username], check_rc=True)
        self.module.run_command(["usermod", "-a", "-G", "HwHiAiUser", self.username], check_rc=True)
    
    def create_log_dir(self):
        self.module.run_command(["mkdir", "-m", "755", "-p", os.path.dirname(self.log_path)], check_rc=True)
        self.module.run_command(["mkdir", "-m", "750", "-p", self.log_path], check_rc=True)
        self.module.run_command(["chown", "{0}:{0}".format(self.username), self.log_path], check_rc=True)

    def install_pkg(self):
        zip_pkg = self._find_npu_exporter_pkg("Ascend-mindxdl-npu-exporter*{}.zip".format(self.arch))
        self.module.run_command(["unzip", zip_pkg, "-d", os.path.dirname(zip_pkg)], check_rc=True)
        pkg = self._find_npu_exporter_pkg("npu-exporter")
        self.module.run_command(["chmod", "500", pkg], check_rc=True)
        self.module.run_command(["cp", pkg, self.install_path], check_rc=True)
        self.module.run_command(["chown", "{0}:{0}".format(self.username), self.install_path], check_rc=True)
    
    def create_service(self):
        cmd = [self.install_path,
               "-ip={}".format(self.ip_address),
               "-port={}".format(self.port),
               "-logFile={}>/dev/null".format(os.path.join(self.log_path, "npu-exporter.log")),
               "2>&1", "&"]
        content = [
            "[Unit]",
            "Description=Ascend npu exporter",
            "Documentation=hiascend.com",
            "[Service]",
            "ExecStart=/bin/bash -c '{}'".format(" ".join(cmd)),
            "RestartSec=2",
            "KillMode=process",
            "Environment='GOGC=50'",
            "Environment='GOMAXPROCS=2'",
            "Environment='GODEBUG=madvdontneed=1'",
            "Type=forking",
            "User={}".format(self.username),
            "Group={}".format(self.username),
            "[Install]",
            "WantedBy=multi-user.target"
        ]
        with open("/etc/systemd/system/npu-exporter.service", 'w') as f:
            f.write("\n".join(content))
    
    def create_timer(self):
        content = [
            "[Unit]",
            "Description=Timer for NPU Exporter Service",
            "[Timer]",
            "OnBootSec=60s",
            "Unit=npu-exporter.service",
            "[Install]",
            "WantedBy=timers.target"
        ]
        with open("/etc/systemd/system/npu-exporter.timer", 'w') as f:
            f.write("\n".join(content))
    
    def setup_service(self):
        self.module.run_command(["systemctl", "enable", "npu-exporter.timer"], check_rc=True)
        self.module.run_command(["systemctl", "start", "npu-exporter"], check_rc=True)
        self.module.run_command(["systemctl", "start", "npu-exporter.timer"], check_rc=True)

    def run(self):
        try:
            self._get_config()
            self.create_user()
            self.create_log_dir()
            self.install_pkg()
            self.create_service()
            self.create_timer()
            self.setup_service()
        except Exception as e:
            self.messages.append(str(e))
            self._module_failed()
        self._module_success()


def main():
    NpuExporterInstaller().run()


if __name__ == "__main__":
    main()
